#include "population.h"

#include "individual.h"
#include "config.h"
#include "prision.h"
#include "logger.h"

#include <cstdlib>
#include <cassert>
#include <iostream>
#include <string>
#include <cmath>
#include <iomanip>
#include <algorithm>


//------------------------------------------------------------------------------
// Public functions

Population::Population(int p_size, float p_crossover_rate, float p_mutation_rate)
    : m_size(p_size)
    , m_crossover_rate(p_crossover_rate)
    , m_mutation_rate(p_mutation_rate)
    , m_prision(NULL)
    , m_current_generation(0)
{
    // Create the starting population with random individuals.
    init();
}

Population::~Population()
{
    if(m_prision) {
        delete m_prision;
        m_prision = NULL;
    }

    clearPopulation();
    clearIntermediatePopulation();
}

void Population::evolve(int p_generations)
{
    // Open log file.
    openLogs();

    // Run the algorithm for p_generations.
    for(int i = 0; i < p_generations; ++i) {
        //-------------------------
        // Selection Process

        // If first run, skeep selection, everyone into the intermediate
        // ... population.
        if(i == 0) {
            // Move the individuals to the intermediate population.
            m_intermediate_individuals = m_individuals;
            // And clear the current individuals.
            m_individuals.clear();
        }
        else {
            // Select from the individual's pool.
            selectIndividuals();
        }

        //-------------------------
        // Crossover Process
        
        crossIndividuals();

        //-------------------------
        // Mutation Process

        mutateIndividuals();

        //-------------------------
        // Evaluation Process

        evaluateIndividuals();

        // We are finished with one generation. Store the 
        // ... biggest generation we have finished.
        m_current_generation = i;

        //-------------------------
        // Collect Stats

        collectStats(i);
    }

    // And close the log files.
    closeLogs();
}


//------------------------------------------------------------------------------
// Private Functions

void Population::init()
{
    Individual *individual;
    int prision_rewards[4];

    // Create the initial individuals.
    m_individuals.clear();
    m_intermediate_individuals.clear();
    for (int i = 0; i < m_size; ++i) {
        // Create a new random individual specifying the history size
        // ... from the config file.
        individual = new Individual();
        individual->randomizeGenes();
        m_individuals.push_back(individual);
    }

    // Create the playground.
    prision_rewards[Prision::S] = PRISION_S;
    prision_rewards[Prision::D] = PRISION_D;
    prision_rewards[Prision::C] = PRISION_C;
    prision_rewards[Prision::T] = PRISION_T;
    m_prision = new Prision(prision_rewards);
}

void Population::selectIndividuals()
{
    // Start the selection process with a clean intermediate
    // ... population.
    clearIntermediatePopulation();

    // Do a roulette selection.
    rouletteSelection();
}

void Population::crossIndividuals()
{
    // Number of individuals in the intermediate population.
    int intermediate_size;
    int parent_a_index;
    int parent_b_index;
    Individual *parent_a;
    Individual *parent_b;
    Individual *child_a;
    Individual *child_b;

    // Work over the intermediate population.

    // Clear the main population first. We are filling it in.
    clearPopulation();

    // Mate individuals until we end up having the desired population
    // ... size.
    intermediate_size = m_intermediate_individuals.size();

    // We will assume that the desired population size is always even.
    assert(m_size % 2 == 0);

    for (int i = 0; i < m_size; i += 2) {
        // Get the random parent indexes to mate.
        parent_a_index = rand() % intermediate_size;
        parent_b_index = rand() % intermediate_size;

        // Get those parents.
        parent_a = m_intermediate_individuals.at(parent_a_index);
        parent_b = m_intermediate_individuals.at(parent_b_index);

        // Will they mate?
        if(std::rand() < m_crossover_rate * RAND_MAX) {
            // Yes...
            // Let them do their thing.
            parent_a->mate(parent_b, child_a, child_b);
            // And add the children to the intermediate population.
            m_individuals.push_back(child_a);
            m_individuals.push_back(child_b);
        }
        else {
            // No... so only clone the parents to the next population.
            m_individuals.push_back(parent_a->clone());
            m_individuals.push_back(parent_b->clone());
        }
    }

    // Done with the intermediate population 
    // ... (born to hate memory leaks!).
    clearIntermediatePopulation();
}

void Population::mutateIndividuals()
{
    std::vector<Individual *>::iterator it;

    // Check if a mutation is happening for each individual.
    for(it = m_individuals.begin(); it != m_individuals.end(); ++it) {
        if(std::rand() < m_mutation_rate * RAND_MAX) {
            // Ahhh! A mutant appears!
            (*it)->mutate();
        }
    }
}

void Population::evaluateIndividuals()
{
    int individual_performance;
    int fitness;

    // Evaluate each individual independently.
    for(int i = 0; i < m_size; ++i) {
        individual_performance = m_prision->pit(i, m_individuals);
        // Let the fitness be the squared performance.
        fitness = individual_performance;
        m_individuals.at(i)->setFitness(fitness);
    }
}

void Population::rouletteSelection()
{
    float average_fitness = 0.0f;
    int copies_to_make;
    float probability_of_extra;
    
    // Get the average fitness of the population.
    average_fitness = getAverageFitness();
    assert((average_fitness != 0.0f));

    // Copy individuals into the intermediate population.
    
    assert(m_intermediate_individuals.size() == 0);

    for(std::vector<Individual *>::iterator it = m_individuals.begin(); 
        it != m_individuals.end(); ++it) {
        copies_to_make = (*it)->getFitness() / average_fitness;
        
        // Add a fixed quantity of clones for this individual.
        for (int i = 0; i < copies_to_make; ++i) {
            // Place a clone in the intermediate population.
            m_intermediate_individuals.push_back((*it)->clone());
        }

        // Should we add more?
        probability_of_extra = (static_cast<float>((*it)->getFitness()) / average_fitness) - 
            static_cast<float>(copies_to_make);

        if(std::rand() < probability_of_extra * RAND_MAX) {
            // Lucky you! Add one more.
            m_intermediate_individuals.push_back((*it)->clone());
        }
    }
    
    
}

void Population::clearPopulation()
{
    std::vector<Individual *>::iterator it;
    for(it = m_individuals.begin(); it != m_individuals.end(); ++it) {
        delete (*it);
    }
    m_individuals.clear();
}

void Population::clearIntermediatePopulation()
{
    std::vector<Individual *>::iterator it;
    for(it = m_intermediate_individuals.begin(); 
        it != m_intermediate_individuals.end(); ++it) {
        delete (*it);
    }
    m_intermediate_individuals.clear();
}

void Population::openLogs()
{
    m_logger.open("average_fitness.dat");
    m_logger.open("max_fitness.dat");
    m_logger.open("min_fitness.dat");
    m_logger.open("bit_count.dat");
    m_logger.open("fittest_individual.dat");
    m_logger.open("mutant.dat");
}

void Population::closeLogs()
{
    m_logger.close();
}

void Population::collectStats(unsigned int p_generation)
{
    Individual *fitest_guy = getFitestIndividual();
    Individual *worst_guy = getWorstIndividual();
    std::string genes = fitest_guy->getGenesAsString();
    unsigned int history_size = fitest_guy->getHistorySize();
    unsigned int string_start = genes.length() - std::pow(2, history_size);
    // Define how many generations to skip to get an average.
    static const int average_cycle = 125;
    // Place to store the average of different variables.
    static int average_average_fitness[2] = {0,125};
    static int average_max_fitness[2] = {0,125};
    static int average_min_fitness[2] = {0,125};

    // Add the variables that we want into the average arrays.
    average_average_fitness[0] += getAverageFitness();
    average_max_fitness[0] += fitest_guy->getFitness();
    average_min_fitness[0] += worst_guy->getFitness();

    if(average_average_fitness[1] == 1) {
        m_logger.log("average_fitness.dat") << std::setw(4)
                                            << p_generation - average_cycle + 1
                                            << std::setw(10)
                                            << average_average_fitness[0] / average_cycle
                                            << std::endl;
        average_average_fitness[0] = 0;
        average_average_fitness[1] = average_cycle;
    } 
    else {
        --average_average_fitness[1];
    }

    if(average_max_fitness[1] == 1) {
        m_logger.log("max_fitness.dat") << std::setw(4)
                                        << p_generation - average_cycle + 1
                                        << std::setw(10)
                                        << average_max_fitness[0] / average_cycle
                                        << std::endl;
        average_max_fitness[0] = 0;
        average_max_fitness[1] = average_cycle;
    }
    else {
        --average_max_fitness[1];
    }

    if(average_min_fitness[1] == 1) {
        m_logger.log("min_fitness.dat") << std::setw(4)
                                        << p_generation - average_cycle + 1
                                        << std::setw(10)
                                        << average_min_fitness[0] / average_cycle
                                        << std::endl;
        average_min_fitness[0] = 0;
        average_min_fitness[1] = average_cycle;
    }
    else {
        --average_min_fitness[1];
    }

    m_logger.log("fittest_individual.dat") << std::setw(4)
                                           << p_generation
                                           << std::setw(8)
                                           << genes.substr(0, I_HISTORY_BITS)
                                           << std::setw(60)
                                           << genes.substr(string_start)
                                           << std::setw(6)
                                           << fitest_guy->getFitness()
                                           << std::endl;

    m_logger.log("bit_count.dat") << std::setw(4)
                                  << p_generation
                                  << std::setw(10)
                                  << std::count(genes.begin(), genes.end(), '1')
                                  << std::endl;

    m_logger.log("mutant.dat") << std::setw(4)
                               << p_generation 
                               << std::setw(6)
                               << fitest_guy->isMutant()
                               << std::endl; 


    //m_logger.log("bit_count.dat") << p_generation
    //                              << std::setw(10)
    //                              << 



    // std::cout << "*----------------------------------------------" 
    //           << std::endl
    //           << "** generation  " << p_generation << std::endl
    //           << "** average fitness  " << getAverageFitness() << std::endl
    //           << "** fitest individual's fitness  " << fitest_guy->getFitness() << std::endl
    //           << "** fitest individual's genes  " << genes.substr(string_start) << std::endl
    //           << "***********************************************" << std::endl;
}

unsigned int Population::getAverageFitness()
{    
    unsigned int average_fitness = 0;
    std::vector<Individual *>::iterator it;

    // Find the average fitness of the population.
    for(it = m_individuals.begin(); it != m_individuals.end(); ++it) {
        average_fitness += (*it)->getFitness();
    }
    average_fitness /= m_individuals.size();
    
    return average_fitness;
}

Individual *Population::getFitestIndividual()
{
    Individual *fitest_individual;
    std::vector<Individual *>::iterator it;

    // Start with the first individual's fitness as the highest.
    it = m_individuals.begin();
    fitest_individual = (*it);
    ++it;

    for(; it != m_individuals.end(); ++it) {
        if(fitest_individual->getFitness() < (*it)->getFitness()) {
            fitest_individual = (*it);
        }
    }

    return fitest_individual;
}

Individual *Population::getWorstIndividual()
{
    Individual *worst_individual;
    std::vector<Individual *>::iterator it;

    // Start with the first individual's fitness as the highest.
    it = m_individuals.begin();
    worst_individual = (*it);
    ++it;

    for(; it != m_individuals.end(); ++it) {
        if(worst_individual->getFitness() > (*it)->getFitness()) {
            worst_individual = (*it);
        }
    }

    return worst_individual;
}
